/*
This is part of Screenbar, which is based in part on the Litestep
shell source code.

Copyright (C) 2008-2009 Jim Babcock
Copyright (C) 1997-2009 The LiteStep Development Team

This program is free software; you can redistribute it and/or
modify it under the terms of the GNU General Public License
as published by the Free Software Foundation; either version 2
of the License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
*/
#include "stdafx.h"
#include <psapi.h>

ProcessTracker *processTracker = NULL;
BOOL SetPrivilege(HANDLE hToken, LPCTSTR lpszPrivilege, BOOL bEnablePrivilege);

/////////////////////////////////////////////////////////////////////////////

ProcessTracker::ProcessTracker()
{
	// Request a privelege token to be able to query other processes' CPU and
	// memory usage.
	HANDLE token;
	if(OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY|TOKEN_ADJUST_PRIVILEGES, &token))
	{
		if(!SetPrivilege(token, SE_DEBUG_NAME, TRUE))
			trace << "Failed to acquire privilege for process tracking\n"; //DEBUG
	}
	else
		trace << "Failed to open process token\n";
	
	CloseHandle(token);
	
	// Query the number of processes
	SYSTEM_INFO systemInfo;
	GetSystemInfo(&systemInfo);
	this->numProcessors = systemInfo.dwNumberOfProcessors;
}

ProcessTracker::~ProcessTracker()
{
	for(map<DWORD,ProcessData*>::iterator ii=processes.begin(); ii!=processes.end(); ii++)
	{
		delete ii->second;
	}
}

// Source: http://msdn.microsoft.com/en-us/library/aa446619(VS.85).aspx
BOOL SetPrivilege(HANDLE hToken, LPCTSTR lpszPrivilege, BOOL bEnablePrivilege)
{
	TOKEN_PRIVILEGES tp;
	LUID luid;

	if(!LookupPrivilegeValue(
			NULL,            // lookup privilege on local system
			lpszPrivilege,   // privilege to lookup
			&luid))        // receives LUID of privilege
	{
		warn(retprintf("LookupPrivilegeValue error: %u\n", GetLastError()).c_str());
		return FALSE;
	}

	tp.PrivilegeCount = 1;
	tp.Privileges[0].Luid = luid;
	if (bEnablePrivilege)
		tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
	else
		tp.Privileges[0].Attributes = 0;

	// Enable the privilege or disable all privileges.

	if(!AdjustTokenPrivileges(
		   hToken,
		   FALSE,
		   &tp,
		   sizeof(TOKEN_PRIVILEGES),
		   (PTOKEN_PRIVILEGES)NULL,
		   (PDWORD)NULL) )
	{
		  warn(retprintf("AdjustTokenPrivileges error: %u\n", GetLastError()).c_str());
		  return FALSE;
	}

	if(GetLastError() == ERROR_NOT_ALL_ASSIGNED)
	{
		  warn(retprintf("The token does not have the specified privilege.\n").c_str());
		  return FALSE;
	}

	return TRUE;
}


void ProcessTracker::update()
{
	// Get a list of PIDs of running processes
	// (This is a whole lot more complicated than it ought to be, but the only
	// way to find out how many processes there are with winapi is to retrieve
	// into a fixed-size array and find out whether or not it was enough.)
	DWORD *pids;
	int procsAlloc = 256;
	int procsReturned = 0;
	
	while(1)
	{
		pids = new DWORD[procsAlloc];
		DWORD bytesReturned = 0;
		
		EnumProcesses(pids, procsAlloc*sizeof(DWORD), &bytesReturned);
		procsReturned = bytesReturned / sizeof(DWORD);
		
		if(procsReturned >= procsAlloc) {
			delete[] pids;
			procsAlloc *= 2;
			pids = new DWORD[procsAlloc];
		} else {
			break;
		}
	}
	
	for(map<DWORD, ProcessData*>::iterator ii=processes.begin(); ii!=processes.end(); ii++)
		ii->second->visited = false;
	
	for(int ii=0; ii<procsReturned; ii++)
	{
		DWORD pid = pids[ii];
		if(pid==0) continue;
		
		map<DWORD, ProcessData*>::iterator it = processes.find(pid);
		
		if(it == processes.end()) {
			addProcess(pid);
			processes[pid]->visited = true;
		} else {
			it->second->visited = true;
			updateProcess(it->second);
		}
	}
	
	for(map<DWORD, ProcessData*>::iterator ii=processes.begin(); ii!=processes.end(); )
	{
		if(ii->second->visited) {
			ii++;
		} else {
			removeProcess(ii->second);
			ii = processes.erase(ii);
		}
	}
	
	delete[] pids;
}

void ProcessTracker::addProcess(DWORD pid)
{
	ProcessData *process = new ProcessData(pid);
	processes[pid] = process;

	HANDLE handle = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, process->pid);
	if(handle)
	{
		PROCESS_MEMORY_COUNTERS memoryCounters;
		memoryCounters.cb = sizeof(memoryCounters);
		GetProcessMemoryInfo(handle, &memoryCounters, sizeof(memoryCounters));
		process->setMemUsage(memoryCounters.WorkingSetSize);
		
		FILETIME creationTime;
		FILETIME exitTime;
		FILETIME kernelTime;
		FILETIME userTime;
		if(GetProcessTimes(handle, &creationTime, &exitTime, &kernelTime, &userTime))
			process->addTimeSample(0, Time64fromFileTime(creationTime));
		else
			trace << "GetProcessTimes failed for pid "<<(int)pid<<"\n"; //DEBUG
	}
	else
	{
		string err = getWindowsErr();
		trace << "OpenProcess failed on PID "<<(int)process->pid<<": "<<err; //DEBUG
	}
	
	CloseHandle(handle);
}

ProcessData *ProcessTracker::getProcess(DWORD pid)
{
	if(processes.find(pid) == processes.end())
		return NULL;
	else
		return processes[pid];
}

void ProcessTracker::updateProcess(ProcessData *process)
{
	HANDLE handle = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, process->pid);
	if(!handle) {
		return;
	}
	
	PROCESS_MEMORY_COUNTERS memoryCounters;
	memoryCounters.cb = sizeof(memoryCounters);
	GetProcessMemoryInfo(handle, &memoryCounters, sizeof(memoryCounters));
	process->setMemUsage(memoryCounters.WorkingSetSize);
	
	FILETIME creationTime;
	FILETIME exitTime;
	FILETIME kernelTime;
	FILETIME userTime;
	if(GetProcessTimes(handle, &creationTime, &exitTime, &kernelTime, &userTime))
	{
		UINT64 currentTime = getSystemTime();
		UINT64 totalTime = Time64fromFileTime(userTime) + Time64fromFileTime(kernelTime);
		
		process->addTimeSample(totalTime, currentTime);
	}
	else trace << "GetProcessTimes failed for pid "<<(int)process->pid<<"\n"; //DEBUG
	
	CloseHandle(handle);
}

void ProcessTracker::removeProcess(ProcessData *process)
{
	delete process;
}

void ProcessTracker::getTaskProcesses(vector<ProcessData*> *outProcesses, TaskData *task)
{
	// TODO
}

void ProcessTracker::getNonTaskProcesses(vector<ProcessData*> *outProcesses, TaskData *task)
{
	// TODO
}

int ProcessTracker::getNumProcessors()
{
	return numProcessors;
}

/////////////////////////////////////////////////////////////////////////////

ProcessData::ProcessData(DWORD pid)
{
	this->pid = pid;
	this->numTimeSamples = 0;
}

void ProcessData::addTimeSample(UINT64 timeUsed, UINT64 timeRecorded)
{
	timeSamples[numTimeSamples%PROCESSDATA_NUM_TIMESAMPLES].used = timeUsed;
	timeSamples[numTimeSamples%PROCESSDATA_NUM_TIMESAMPLES].recorded = timeRecorded;
	numTimeSamples++;
}

float ProcessData::getCPUUsage()
{
	if(numTimeSamples < 2)
		return 0;
	
	TimeSample *newestSample;
	TimeSample *oldestSample;
	int pos = numTimeSamples%PROCESSDATA_NUM_TIMESAMPLES;
	
	if(pos==0)
		newestSample = &timeSamples[PROCESSDATA_NUM_TIMESAMPLES-1];
	else
		newestSample = &timeSamples[pos-1];
	
	if(pos==numTimeSamples)
		oldestSample = &timeSamples[0];
	else
		oldestSample = &timeSamples[pos];
	
	UINT64 interval = newestSample->recorded - oldestSample->recorded;
	UINT64 usage = newestSample->used - oldestSample->used;
	
	assert(newestSample->recorded >= oldestSample->recorded);
	assert(newestSample->used >= oldestSample->used);
	assert(interval > 0);
	
	return (float)usage / ((float)interval * processTracker->getNumProcessors());
}

void ProcessData::setMemUsage(size_t bytesUsed)
{
	this->memUsage = bytesUsed;
}

size_t ProcessData::getMemUsage()
{
	return memUsage;
}

/////////////////////////////////////////////////////////////////////////////
// Process parent identification. There is no published Windows API for this;
// instead, use NtQueryInformationProcess, which is considered "internal".
#include <winternl.h>

// This is the same as the PROCESS_BASIC_INFORMATION structure defined in
// <winternl.h>, except that the reserved fields have been given names and
// types, from <http://www.daniweb.com/forums/thread114975.html>.
typedef struct _PROCESS_BASIC_INFORMATION_2
{
    NTSTATUS ExitStatus;
    PPEB PebBaseAddress;
    unsigned long AffinityMask;
    long BasePriority;
    unsigned long UniqueProcessId;
    unsigned long InheritedFromUniqueProcessId;
} PBI;

typedef LONG (__stdcall *QueryInformationProcessFunc)
	(HANDLE, PROCESSINFOCLASS, PVOID, ULONG, PULONG);

// This function is cleaned up and adapted from an article by J|rgen Kraus,
// <http://www.codeguru.com/Cpp/W-P/win32/article.php/c1437/>.
DWORD GetParentProcessID(DWORD pid)
{
	DWORD parentPID = 0;

	//  create entry point for 'NtQueryInformationProcess()'
	HMODULE ntDll = GetModuleHandleA("ntdll");
	QueryInformationProcessFunc NtQueryInformationProcess =
		(QueryInformationProcessFunc)GetProcAddress(ntDll, "NtQueryInformationProcess");
	
	HANDLE hProcess = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, pid);

	// could fail due to invalid PID or insufficiant privileges
	if(!hProcess)
		return -1;

	// gather information
	PBI processInfo;
	ULONG retLen;

	if(!NtQueryInformationProcess(
		hProcess, ProcessBasicInformation, (void*)&processInfo, sizeof(processInfo), &retLen))
	{
		parentPID = processInfo.InheritedFromUniqueProcessId;
	}
	
	CloseHandle(hProcess);

	return parentPID;
}

